from common.variables import get_sqlalchemy_url, TimeFrameType
from sqlalchemy import Column, create_engine, DateTime, Float, JSON, String
from sqlalchemy.orm import declarative_base, sessionmaker


Base = declarative_base()


class KlineDataBase:
    __tablename__: str
    symbol = Column(String, primary_key=True)  # 股票代码
    xdatetime = Column(DateTime, primary_key=True)  # 时间戳
    fname = Column(String, primary_key=True)  # 因子名称
    value = Column(Float)  # 因子值

    @classmethod
    def configure_table_args(cls, database_dialect):
        """
        根据数据库类型配置分区设置，支持 MySQL, ClickHouse 和 DuckDB
        """
        if not hasattr(cls, "__table_args__"):
            cls.__table_args__ = {}
        if database_dialect.startswith("mysql"):
            cls.__table_args__.update({"mysql_partition_by": ["date(xdatetime)"]})
        elif database_dialect.startswith("clickhouse"):
            cls.__table_args__.update(
                {
                    "clickhouse_engine": "ReplacingMergeTree()",
                    "clickhouse_partition_by": ["toYYYYMM(xdatetime)"],
                }
            )
        elif database_dialect.startswith("duckdb"):
            cls.__table_args__.update({"duckdb_partition_by": ["date(xdatetime)"]})


class KlineData1m(KlineDataBase, Base):
    __tablename__ = "kline_data_1m"


class KlineData5(KlineDataBase, Base):
    __tablename__ = "kline_data_5m"


class KlineData10m(KlineDataBase, Base):
    __tablename__ = "kline_data_10m"


for orm in [KlineData1m, KlineData5, KlineData10m]:
    orm.configure_table_args(get_sqlalchemy_url("db_test3"))


class ExperimentDB(Base):
    __tablename__: str = "experi_tb"
    expid = Column(String, primary_key=True)
    title = Column(String, doc="自定义名称")
    name = Column(String, doc="任务类型")
    create = Column(DateTime, doc="创建时间")
    start = Column(DateTime, doc="开始时间")
    end = Column(DateTime, doc="结束时间")
    status = Column(String, doc="状态")
    info = Column(JSON, doc="包括参数等信息")
    results = Column(JSON, doc="结果")


def get_session_factory(db_name: str):
    engine = create_engine(get_sqlalchemy_url(db_name), echo=False)
    Base.metadata.create_all(engine)
    Session = sessionmaker(bind=engine)
    return Session


def get_session():
    Session = get_session_factory(db_name="db_test3")
    session = Session()
    return session


def get_kline_table_orm(timeframe: TimeFrameType):
    if timeframe == "1m":
        return KlineData1m
    elif timeframe == "5m":
        return KlineData5
    elif timeframe == "10m":
        return KlineData10m
    else:
        raise ValueError(f"timeframe {timeframe} not supported")
